import tsMorph = require("ts-morph");

/**
 * 获取函数的所有泛型参数名称
 * @param sourceCode 函数所在的源代码
 * @param functionName 目标函数名
 * @returns 泛型名称数组，如 ["T", "U"]
 */
function getFunctionGenericNames(
  sourceCode: string,
  functionName: string
): string[] {
  const project = new tsMorph.Project({ useInMemoryFileSystem: true });
  const sourceFile = project.createSourceFile("temp.ts", sourceCode);

  // 查找函数声明
  const functionDecl = sourceFile.getFunction(functionName);
  if (functionDecl) {
    return extractGenericsFromFunction(functionDecl);
  }

  // 查找箭头函数或函数表达式（通过变量声明）
  const variableDecl = sourceFile.getVariableDeclaration(functionName);
  if (variableDecl && tsMorph.Node.isFunctionExpression(variableDecl.getInitializer())) {
    const funcExpr = variableDecl.getInitializer()!;
    return extractGenericsFromFunction(funcExpr);
  }

  throw new Error(`Function ${functionName} not found`);
}

// 从函数节点中提取泛型
function extractGenericsFromFunction(
  func: tsMorph.FunctionDeclaration | tsMorph.Node
): string[] {

  if (tsMorph.Node.isFunctionExpression(func) || tsMorph.Node.isArrowFunction(func)) {
    const typeParams = func.getTypeParameters();
    return typeParams.map(tp => tp.getName());
  }

  if (tsMorph.Node.isFunctionDeclaration(func)) {
    return func.getTypeParameters().map(tp => tp.getName());
  }
  return [];
}

export = getFunctionGenericNames;
